#include "point_GS.hpp"
#include "scal_kernels.hpp"

template<typename idx_t, typename data_t, typename setup_t, typename calc_t, int dof>
PointGS<idx_t, data_t, setup_t, calc_t, dof>::PointGS() : Solver<idx_t, setup_t, calc_t, dof>() {}

template<typename idx_t, typename data_t, typename setup_t, typename calc_t, int dof>
PointGS<idx_t, data_t, setup_t, calc_t, dof>::PointGS(SCAN_TYPE type) : Solver<idx_t, setup_t, calc_t, dof>()
{
    scan_type = type;
    if (type == FORW_BACK)      last_time_forward = false;
    else if (type == BACK_FORW) last_time_forward = true;
}

template<typename idx_t, typename data_t, typename setup_t, typename calc_t, int dof>
PointGS<idx_t, data_t, setup_t, calc_t, dof>::~PointGS()
{
    if (owns_offd) { delete pgs_offd; pgs_offd = nullptr; }
    delete [] offd_x;
    delete [] send_buf;
    delete [] tmp_rhs;
    delete [] map_new2old; delete [] map_old2new;
}

template<typename idx_t, typename data_t, typename setup_t, typename calc_t, int dof>
void PointGS<idx_t, data_t, setup_t, calc_t, dof>::SetScanType(SCAN_TYPE type)
{
    scan_type = type;
}

template<typename idx_t, typename data_t, typename setup_t, typename calc_t, int dof>
void PointGS<idx_t, data_t, setup_t, calc_t, dof>::truncate() {}


template<typename idx_t, typename data_t, typename setup_t, typename calc_t, int dof>
void PointGS<idx_t, data_t, setup_t, calc_t, dof>::SetOperator(const Operator<idx_t, setup_t, dof> & op)
{
    const par_CSRMatrix<idx_t, setup_t, setup_t, dof> & par_A = (const par_CSRMatrix<idx_t, setup_t, setup_t, dof> &) op;
    int my_pid; MPI_Comm_rank(par_A.comm, & my_pid);
    if (my_pid == 0) printf("Setup PGS #dof %d weight %.3f\n", dof, this->weight);
    assert(this->setup_called == false);
    this->setup_time = - wall_time();

    if (false) {
        // this->scan_type = FORWARD;

        map_old2new = new idx_t [par_A.diag.nrows];
        map_new2old = new idx_t [par_A.diag.nrows];

        par_CSRMatrix<idx_t, setup_t, setup_t, 1> S_mat(MPI_COMM_WORLD, par_A.rows_partition, par_A.cols_partition);
        S_mat.col_map_offd      = par_A.col_map_offd;
        S_mat.commpkg           = par_A.commpkg;
        S_mat.commpkg_transpose = par_A.commpkg_transpose;

        S_mat.diag.nrows = par_A.diag.nrows;
        S_mat.diag.ncols = par_A.diag.ncols;
        S_mat.diag.nnz   = par_A.diag.nnz  ; S_mat.diag.alloc_mem(false);
        S_mat.offd.nrows = par_A.offd.nrows;
        S_mat.offd.ncols = par_A.offd.ncols;
        S_mat.offd.nnz   = par_A.offd.nnz  ; S_mat.offd.alloc_mem(false);

        constexpr int e_size = dof*dof;
        auto calc_Frob = [] (const setup_t * val) {
            setup_t ret = 0.0;
            for (int f = 0; f < e_size; f++) ret += val[f] * val[f];
            return sqrt(ret);
        };

        for (idx_t i = 0; i <= S_mat.diag.nrows; i++)
            S_mat.diag.row_ptr[i] = par_A.diag.row_ptr[i];

        for (idx_t p = 0; p <  S_mat.diag.row_ptr[S_mat.diag.nrows]; p ++) {
            S_mat.diag.col_idx[p] = par_A.diag.col_idx[p];
            S_mat.diag.vals   [p] = calc_Frob(par_A.diag.vals + p * e_size);
        }

        for (idx_t i = 0; i <= S_mat.offd.nrows; i++)
            S_mat.offd.row_ptr[i] = par_A.offd.row_ptr[i];

        for (idx_t p = 0; p <  S_mat.offd.row_ptr[S_mat.offd.nrows]; p++) {
            S_mat.offd.col_idx[p] = par_A.offd.col_idx[p];
            S_mat.offd.vals   [p] = calc_Frob(par_A.offd.vals + p * e_size);
        }

        struct Handle {
            idx_t id = -1;
            idx_t cnt = 0;
            setup_t wgt = 0.0;
            bool operator < (const Handle & b) const {
                // if (cnt == b.cnt) return id < b.id;
                // else return cnt < b.cnt;
                
                return wgt < b.wgt;// 受别人影响小的排在前面
            }
        };
        Handle * ndegs = new Handle [par_A.diag.nrows];

        // if (my_pid == 0) printf("Calc done\n");

        for (idx_t i = 0; i < S_mat.diag.nrows; i++) {
            idx_t row_nnz = S_mat.diag.row_ptr[i + 1] - S_mat.diag.row_ptr[i]
                        +   S_mat.offd.row_ptr[i + 1] - S_mat.offd.row_ptr[i];
            idx_t ngb_ids [row_nnz - 1], cnt = 0;
            setup_t strengths[row_nnz - 1], diag_strength = 0.0;
            for (idx_t p = S_mat.diag.row_ptr[i]; p < S_mat.diag.row_ptr[i + 1]; p++) {
                if (S_mat.diag.col_idx[p] == i) {
                    diag_strength = S_mat.diag.vals[p];
                } else {
                    ngb_ids[cnt] = S_mat.beg_col + S_mat.diag.col_idx[p];
                    strengths[cnt] = S_mat.diag.vals[p];
                    cnt ++;
                }
            }
            for (idx_t p = S_mat.offd.row_ptr[i]; p < S_mat.offd.row_ptr[i + 1]; p++) {
                ngb_ids[cnt] = S_mat.col_map_offd[S_mat.offd.col_idx[p]];
                strengths[cnt] = S_mat.offd.vals[p];
                cnt ++;
            }
            assert(diag_strength != 0.0);
            assert(cnt == row_nnz - 1);
            for (idx_t j = 0; j < cnt; j++) strengths[j] /= diag_strength;
            // 按强度降序排列
            sort_vals_with_indices<idx_t, setup_t>(cnt, strengths, ngb_ids);

            ndegs[i].id  = i;
            ndegs[i].wgt = strengths[0];
            const setup_t threshold = 1.0 / 6.0;
            for (idx_t j = 0; j < cnt; j++) {
                if (strengths[j] < threshold) { break; }
                else ndegs[i].cnt ++;
            }
        }
        // 排序
        std::sort(ndegs, ndegs + par_A.diag.nrows);
        // 记录
        for (idx_t new_i = 0; new_i < S_mat.diag.nrows; new_i ++) {
            map_new2old[new_i] = ndegs[new_i].id;
            map_old2new[ndegs[new_i].id] = new_i;
        }

        delete [] ndegs;

        // 构造重排后的局部矩阵
        CSRMatrix<idx_t, setup_t, setup_t, dof> new_diag(par_A.diag.nrows, par_A.diag.ncols, par_A.diag.nnz);
        new_diag.row_ptr[0] = 0;
        for (idx_t new_i = 0; new_i < new_diag.nrows; new_i ++) {
            const idx_t old_i = map_new2old[new_i];
            new_diag.row_ptr[new_i + 1] = par_A.diag.row_ptr[old_i + 1] - par_A.diag.row_ptr[old_i];
        }
        for (idx_t new_i = 0; new_i < new_diag.nrows; new_i ++) {
            new_diag.row_ptr[new_i + 1] += new_diag.row_ptr[new_i];
        }
        assert(new_diag.row_ptr[new_diag.nrows] == par_A.diag.row_ptr[new_diag.nrows]);

        for (idx_t new_i = 0; new_i < new_diag.nrows; new_i ++) {
            const idx_t old_i = map_new2old[new_i];
            const idx_t old_offset = par_A.diag.row_ptr[old_i],
                        new_offset =   new_diag.row_ptr[new_i];
            const idx_t num = new_diag.row_ptr[new_i + 1] - new_offset;
            assert(num == par_A.diag.row_ptr[old_i + 1] - old_offset);
            for (idx_t p = 0; p < num; p++)
                new_diag.col_idx[new_offset + p] = map_old2new[ par_A.diag.col_idx[old_offset + p] ];// 注意将列序号转为重排后的
            for (idx_t p = 0; p < num * e_size; p++)
                new_diag.vals[new_offset * e_size + p] = par_A.diag.vals[old_offset * e_size + p];

            sort_indices_with_vals<idx_t, setup_t, dof>(num,    new_diag.col_idx + new_offset,
                                                                new_diag.vals    + new_offset * e_size);
        }

        separate_LU(new_diag);
    } else {
        separate_LU(par_A.diag);
    }

    // 非对角部分，拷贝或借用
    A_calc = & par_A;
    if constexpr (sizeof(setup_t) == sizeof(data_t) && sizeof(setup_t) == sizeof(calc_t)) {
        pgs_offd = & par_A.offd;
    } else {
        owns_offd = true;
        CSRMatrix<idx_t, data_t, calc_t, dof> * cpy_offd = new CSRMatrix<idx_t, data_t, calc_t, dof>;
        cpy_offd->nrows = par_A.offd.nrows;
        cpy_offd->ncols = par_A.offd.ncols;
        cpy_offd->nnz   = par_A.offd.nnz;
        cpy_offd->alloc_mem(false);

        #pragma omp parallel
        {
            #pragma omp for schedule(static) nowait
            for (idx_t i = 0; i <= cpy_offd->nrows; i++)
                cpy_offd->row_ptr[i] = par_A.offd.row_ptr[i];
            
            #pragma omp for schedule(static) nowait
            for (idx_t j = 0; j < cpy_offd->row_ptr[cpy_offd->nrows]; j++)
                cpy_offd->col_idx[j] = par_A.offd.col_idx[j];
            
            #pragma omp for schedule(static) nowait
            for (long long j = 0; j < (long long) cpy_offd->row_ptr[cpy_offd->nrows] * dof*dof; j++)
                cpy_offd->vals   [j] = par_A.offd.vals   [j];
        }

        pgs_offd = cpy_offd;
    }
    offd_x   = new calc_t [pgs_offd->ncols * dof];
    send_buf = new calc_t [par_A.commpkg->send_map_starts[par_A.commpkg->num_sends] * dof];
    tmp_rhs  = new calc_t [pgs_offd->nrows * dof];
    
    this->setup_called = true;
    this->setup_time += wall_time();
}

template<typename idx_t, typename data_t, typename setup_t, typename calc_t, int dof>
void PointGS<idx_t, data_t, setup_t, calc_t, dof>::Mult_ipl(const par_Vector<idx_t, calc_t, dof> & b,
    par_Vector<idx_t, calc_t, dof> & x) const
{
    // ===============================
    //          Forward Pass
    // ===============================
    if (scan_type == SYMMETRIC || scan_type == FORWARD) {
        if (this->zero_guess == false) {// 需要做一次非对角部分的spmv，修改右端项b，故需通信非结构部分非零元对应的向量
            A_calc->commpkg->init_update_halo(dof, x.local_vector->data, send_buf, offd_x);
            A_calc->commpkg->wait_update_halo();

            // double ck_dot = 0.0;
            // for (idx_t i = 0; i < pgs_offd->ncols; i++) ck_dot += offd_x[i] * offd_x[i];
            // printf("  F  offd_x ck %.5e\n", ck_dot);

            pgs_offd->Mult(-1.0, offd_x, 1.0, b.local_vector->data, tmp_rhs);// tmpRHS <- b - A_offd*x_offd
            wrapper_forward(      tmp_rhs       , x.local_vector->data);
        } else {
            x.set_val(0.0);
            wrapper_forward(b.local_vector->data, x.local_vector->data);
        }
        this->zero_guess = false;
    }


    // ===============================
    //         Backward Pass
    // ===============================
    if (scan_type == SYMMETRIC || scan_type == BACKWARD) {
        if (this->zero_guess == false) {// 需要做一次非对角部分的spmv，修改右端项b，故需通信非结构部分非零元对应的向量
            A_calc->commpkg->init_update_halo(dof, x.local_vector->data, send_buf, offd_x);
            A_calc->commpkg->wait_update_halo();

            // printf("bp 000\n");
            // printf("PGS offd nrows %d ncols %d nnz %d x nrows %d\n", pgs_offd->nrows, pgs_offd->ncols, pgs_offd->nnz,
            //     x.local_vector->num);
            // printf("rpt %d %d %d ... %d %d %d\n",   pgs_offd->row_ptr[0], pgs_offd->row_ptr[1], pgs_offd->row_ptr[2],
            //     pgs_offd->row_ptr[pgs_offd->nrows-2], pgs_offd->row_ptr[pgs_offd->nrows-1],
            //     pgs_offd->row_ptr[pgs_offd->nrows  ]);

            // double ck_dot = 0.0;
            // for (idx_t i = 0; i < pgs_offd->ncols; i++) ck_dot += offd_x[i] * offd_x[i];
            // printf("  B  offd_x ck %.5e\n", ck_dot);

            pgs_offd->Mult(-1.0, offd_x, 1.0, b.local_vector->data, tmp_rhs);// tmpRHS <- b - A_offd*x_offd
            // printf("bp 010\n");
            wrapper_backward(      tmp_rhs       , x.local_vector->data);
        } else {
            x.set_val(0.0);
            wrapper_backward(b.local_vector->data, x.local_vector->data);
        }
        this->zero_guess = false;
    }
}

template<typename idx_t, typename data_t, typename setup_t, typename calc_t, int dof>
void PointGS<idx_t, data_t, setup_t, calc_t, dof>::wrapper_forward(const calc_t * rhs, calc_t * sol) const
{
    if (map_new2old) {// 需要重排序
        assert(false);
        const idx_t nrows = sep_L.nrows;
        calc_t * new_rhs = new calc_t [nrows * dof];
        calc_t * new_sol = new calc_t [nrows * dof];

        for (idx_t new_i = 0; new_i < nrows; new_i ++) {
            const idx_t old_i = map_new2old[new_i];
            for (int f = 0; f < dof; f++) {
                new_rhs[new_i * dof + f] = rhs[old_i * dof + f];
                new_sol[new_i * dof + f] = sol[old_i * dof + f];
            }
        }

        forward(new_rhs, new_sol);

        for (idx_t old_i = 0; old_i < nrows; old_i ++) {
            const idx_t new_i = map_old2new[old_i];
            for (int f = 0; f < dof; f++) {
                sol[old_i * dof + f] = new_sol[new_i * dof + f];
            }
        }

        delete [] new_rhs; delete [] new_sol;
    } else {
        forward(rhs, sol);
    }
}

template<typename idx_t, typename data_t, typename setup_t, typename calc_t, int dof>
void PointGS<idx_t, data_t, setup_t, calc_t, dof>::wrapper_backward(const calc_t * rhs, calc_t * sol) const
{
    if (map_new2old) {// 需要重排序
        assert(false);
        const idx_t nrows = sep_L.nrows;
        calc_t * new_rhs = new calc_t [nrows * dof];
        calc_t * new_sol = new calc_t [nrows * dof];

        for (idx_t new_i = 0; new_i < nrows; new_i ++) {
            const idx_t old_i = map_new2old[new_i];
            for (int f = 0; f < dof; f++) {
                new_rhs[new_i * dof + f] = rhs[old_i * dof + f];
                new_sol[new_i * dof + f] = sol[old_i * dof + f];
            }
        }

        backward(new_rhs, new_sol);

        for (idx_t old_i = 0; old_i < nrows; old_i ++) {
            const idx_t new_i = map_old2new[old_i];
            for (int f = 0; f < dof; f++) {
                sol[old_i * dof + f] = new_sol[new_i * dof + f];
            }
        }

        delete [] new_rhs; delete [] new_sol;
    } else {
        backward(rhs, sol);
    }
}

template<typename idx_t, typename data_t, typename setup_t, typename calc_t, int dof>
void PointGS<idx_t, data_t, setup_t, calc_t, dof>::forward(const calc_t * rhs, calc_t * sol) const
{
    const calc_t weight = this->weight;
    const calc_t one_minus_weight = 1.0 - weight;
    const CSRMatrix<idx_t, data_t, calc_t, dof> & L = sep_L, & U = sep_U;
    const idx_t nrows = sep_L.nrows;
    calc_t * old = new calc_t [nrows * dof];

    // double u_dot = 0.0, f_dot = 0.0;
    // for (idx_t i = 0; i < nrows; i++) {
    //     u_dot += sol[i] * sol[i];
    //     f_dot += rhs[i] * rhs[i];
    // }
    // printf("  before forward (u,u)=%.5e (f,f)=%.5e\n", u_dot, f_dot);

    #pragma omp parallel
    {
        // 线程初始化，负载均衡
        int tid = omp_get_thread_num();
        int nt  = omp_get_num_threads();
        idx_t tbeg = 0, tend = 0;
        thread_load_balance(nrows, nt, tid, tbeg, tend);

        // 前扫: D*x^{t+1} = b - L*x^{t+1}
        if (this->zero_guess) {
            // 零初值优化的前扫，只需要计算 “本线程” 对角块的部分
            if constexpr (dof == 1) {
                for (idx_t i = tbeg; i < tend; i++) {
                    calc_t t_res = rhs[i];
                    for (auto p = L.row_ptr[i]; p < L.row_ptr[i + 1]; p++) {
                        const idx_t j = L.col_idx[p]; assert(j < i);
                        if (tbeg <= j) // 本线程范围内
                            t_res -= L.vals[p] * sol[j];
                    }
                    sol[i] = this->weight * t_res * inv_D[i];
                }
            } else { static_assert(dof > 1);
                constexpr int e_size = dof*dof;
                for (idx_t i = tbeg; i < tend; i++) {
                    calc_t t_res [dof];
                    for (int f = 0; f < dof; f++) t_res[f] = - rhs[i * dof + f];// 注意这里是取了负
                    for (auto p = L.row_ptr[i]; p < L.row_ptr[i + 1]; p++) {
                        const idx_t j = L.col_idx[p]; assert(j < i);
                        if (tbeg <= j)// 本线程范围内
                            matvec_mla<idx_t, data_t, calc_t, dof>(L.vals + p * e_size, sol + j * dof, t_res);
                    }
                    // 注意最后再取回负
                    matvec_mul<idx_t, calc_t, calc_t, dof>(& inv_D[i * e_size], t_res, sol + i * dof, - weight);
                }
            }
        }
        // 非零初值优化的前扫
        // 线程内 Gauss-Seidel
        // 线程间 Jacobi
        // 进程间 Jacobi
        else {
            for (idx_t i = tbeg * dof; i < tend * dof; i++)
                old[i] = sol[i];
            
            #pragma omp barrier

            if constexpr (dof == 1) {
                for (idx_t i = tbeg; i < tend; i++) {
                    calc_t t_res = rhs[i];
                    for (auto p = L.row_ptr[i]; p < L.row_ptr[i + 1]; p++) {
                        const idx_t j = L.col_idx[p]; assert(j < tend);// 因为 j < i < tend
                        if (tbeg <= j)  t_res -= L.vals[p] * sol[j];// 本线程范围内，用新值
                        else            t_res -= L.vals[p] * old[j];// 其他线程的范围，用旧值
                    }
                    for (auto p = U.row_ptr[i]; p < U.row_ptr[i + 1]; p++) {// 前扫时上三角部分的必然都是旧值
                        const idx_t j = U.col_idx[p];// j 必然大于 i
                        t_res -= U.vals[p] * old[j];
                    }
                    sol[i] = one_minus_weight * sol[i] + weight * t_res * inv_D[i];
                }
            } else { static_assert(dof > 1);
                constexpr int e_size = dof*dof;
                for (idx_t i = tbeg; i < tend; i++) {
                    calc_t t_res [dof];
                    for (int f = 0; f < dof; f++) t_res[f] = - rhs[i * dof + f];// 注意这里是取了负
                    for (auto p = L.row_ptr[i]; p < L.row_ptr[i + 1]; p++) {
                        const idx_t j = L.col_idx[p]; assert(j < tend);// 因为 j < i < tend
                        if (tbeg <= j)  matvec_mla<idx_t, data_t, calc_t, dof>(L.vals + p * e_size, sol + j * dof, t_res);
                        else            matvec_mla<idx_t, data_t, calc_t, dof>(L.vals + p * e_size, old + j * dof, t_res);
                    }
                    for (auto p = U.row_ptr[i]; p < U.row_ptr[i + 1]; p++) {// 前扫时上三角部分的必然都是旧值
                        const idx_t j = U.col_idx[p];// j 必然大于 i
                        matvec_mla<idx_t, data_t, calc_t, dof>(U.vals + p * e_size, old + j * dof, t_res);
                    }
                    // 注意最后再取回负
                    matvec_mla<idx_t, calc_t, calc_t, dof>(- weight, & inv_D[i * e_size], t_res, one_minus_weight, sol + i * dof, sol + i * dof);
                }
            }
        }// non-zero guess
    }// omp
    delete [] old;

    // u_dot = f_dot = 0.0;
    // for (idx_t i = 0; i < nrows; i++) {
    //     u_dot += sol[i] * sol[i];
    //     f_dot += rhs[i] * rhs[i];
    // }
    // printf("  after  forward (u,u)=%.5e (f,f)=%.5e\n", u_dot, f_dot);
}

template<typename idx_t, typename data_t, typename setup_t, typename calc_t, int dof>
void PointGS<idx_t, data_t, setup_t, calc_t, dof>::backward(const calc_t * rhs, calc_t * sol) const
{
    const calc_t weight = this->weight;
    const calc_t one_minus_weight = 1.0 - weight;
    const CSRMatrix<idx_t, data_t, calc_t, dof> & L = sep_L, & U = sep_U;
    const idx_t nrows = sep_L.nrows;
    calc_t * old = new calc_t [nrows * dof];

    // double u_dot = 0.0, f_dot = 0.0;
    // for (idx_t i = 0; i < nrows; i++) {
    //     u_dot += sol[i] * sol[i];
    //     f_dot += rhs[i] * rhs[i];
    // }
    // printf("  before backward (u,u)=%.5e (f,f)=%.5e\n", u_dot, f_dot);

    #pragma omp parallel
    {
        // 线程初始化，负载均衡
        int tid = omp_get_thread_num();
        int nt  = omp_get_num_threads();
        idx_t tbeg = 0, tend = 0;
        thread_load_balance(nrows, nt, tid, tbeg, tend);

        if (this->zero_guess) {
            assert(false);
        }
        // 后扫，非零初值
        // 线程内 Gauss-Seidel
        // 线程间 Jacobi
        // 进程间 Jacobi
        else {
            for (idx_t i = tbeg * dof; i < tend * dof; i++)
                old[i] = sol[i];
            
            #pragma omp barrier

            if constexpr (dof == 1) {
                for (idx_t i = tend - 1; i >= tbeg; i--) {
                    calc_t t_res = rhs[i];
                    for (auto p = L.row_ptr[i]; p < L.row_ptr[i + 1]; p++) {// 后扫时上三角部分的必然都是旧值
                        const idx_t j = L.col_idx[p]; assert(j < i);// j 必然小于 i
                        t_res -= L.vals[p] * old[j];
                    }
                    for (auto p = U.row_ptr[i]; p < U.row_ptr[i + 1]; p++) {
                        const idx_t j = U.col_idx[p]; assert(i < j);// 注意 tbeg <= i < j
                        if (j < tend)   t_res -= U.vals[p] * sol[j];// 本线程范围内，用新值
                        else            t_res -= U.vals[p] * old[j];// 其他线程的范围，用旧值
                    }
                    sol[i] = one_minus_weight * sol[i] + weight * t_res * inv_D[i];
                }
            } else { static_assert(dof > 1);
                constexpr int e_size = dof*dof;
                for (idx_t i = tend - 1; i >= tbeg; i--) {
                    calc_t t_res [dof];
                    for (int f = 0; f < dof; f++) t_res[f] = - rhs[i * dof + f];// 注意这里是取了负
                    for (auto p = L.row_ptr[i]; p < L.row_ptr[i + 1]; p++) {// 后扫时上三角部分的必然都是旧值
                        const idx_t j = L.col_idx[p];// j 必然小于 i
                        matvec_mla<idx_t, data_t, calc_t, dof>(L.vals + p * e_size, old + j * dof, t_res);
                    }
                    for (auto p = U.row_ptr[i]; p < U.row_ptr[i + 1]; p++) {// 前扫时上三角部分的必然都是旧值
                        const idx_t j = U.col_idx[p]; assert(tbeg <= j);// 因为 tbeg <= i < j
                        if (j < tend)   matvec_mla<idx_t, data_t, calc_t, dof>(U.vals + p * e_size, sol + j * dof, t_res);
                        else            matvec_mla<idx_t, data_t, calc_t, dof>(U.vals + p * e_size, old + j * dof, t_res);
                    }
                    // 注意最后再取回负
                    matvec_mla<idx_t, calc_t, calc_t, dof>(- weight, & inv_D[i * e_size], t_res, one_minus_weight, sol + i * dof, sol + i * dof);
                }
            }
        }// non-zero guess
    }// omp
    delete [] old;

    // u_dot = f_dot = 0.0;
    // for (idx_t i = 0; i < nrows; i++) {
    //     u_dot += sol[i] * sol[i];
    //     f_dot += rhs[i] * rhs[i];
    // }
    // printf("  after  backward (u,u)=%.5e (f,f)=%.5e\n", u_dot, f_dot);
}


template<typename idx_t, typename data_t, typename setup_t, typename calc_t, int dof>
void PointGS<idx_t, data_t, setup_t, calc_t, dof>::separate_LU(const CSRMatrix<idx_t, setup_t, setup_t, dof> & A)
{
    const idx_t nrows = A.nrows;
    CSRMatrix<idx_t, data_t, calc_t, dof> & L = sep_L;
    CSRMatrix<idx_t, data_t, calc_t, dof> & U = sep_U;

    L.nrows = U.nrows = nrows;
    L.row_ptr = new idx_t [nrows + 1]; L.row_ptr[0] = 0; 
    U.row_ptr = new idx_t [nrows + 1]; U.row_ptr[0] = 0;

    for (idx_t i = 0; i < nrows; i++) {
        for (auto p = A.row_ptr[i]; p < A.row_ptr[i + 1] - 1; p++)// 检查每行内是否升序排列
            assert(A.col_idx[p] < A.col_idx[p + 1]);
        
        idx_t   lnz = 0,// 下三角非零元数
                rnz = 0;// 上三角非零元数（注意不包含主对角元！！！）
        for (auto p = A.row_ptr[i]; p < A.row_ptr[i + 1]; p++) {
            if      (A.col_idx[p] <  i) lnz ++;
            else break;
        }
        assert(A.col_idx[A.row_ptr[i] + lnz] == i);
        rnz = A.row_ptr[i + 1] - A.row_ptr[i] - 1 - lnz;// 刨除主对角元
        L.row_ptr[i + 1] = L.row_ptr[i] + lnz;
        U.row_ptr[i + 1] = U.row_ptr[i] + rnz;
    }

    constexpr int e_size = dof*dof;
    L.col_idx = new idx_t [L.row_ptr[nrows]];
    U.col_idx = new idx_t [U.row_ptr[nrows]];
    L.vals    = new data_t[L.row_ptr[nrows] * e_size];
    U.vals    = new data_t[U.row_ptr[nrows] * e_size];
    inv_D.resize(nrows * e_size);// 主对角元的逆矩阵

    for (idx_t i = 0; i < nrows; i++) {
        const idx_t lnz = L.row_ptr[i + 1] - L.row_ptr[i],
                    rnz = U.row_ptr[i + 1] - U.row_ptr[i];
        assert(lnz + 1 + rnz == A.row_ptr[i + 1] - A.row_ptr[i]);
        assert(A.col_idx[A.row_ptr[i] + lnz] == i);
        // 拷贝下三角部分
        for (idx_t j = 0; j < lnz; j++) {
            L.col_idx[L.row_ptr[i] + j] = A.col_idx[A.row_ptr[i] + j];
        }
        for (idx_t j = 0; j < lnz * e_size; j++) {
            L.vals[L.row_ptr[i] * e_size + j] = A.vals[A.row_ptr[i] * e_size + j];
        }
        // 主对角元求逆
        if constexpr (dof == 1) {
            inv_D[i] = 1.0 / A.vals[A.row_ptr[i] + lnz];
        } else { static_assert(dof > 1);
            setup_t src_buf [e_size], inv_buf [e_size];
            for (int f = 0; f < e_size; f++) src_buf[f] = A.vals[(A.row_ptr[i] + lnz) * e_size + f];
            matinv_row<setup_t, dof>(src_buf, inv_buf);
            for (int f = 0; f < e_size; f++) inv_D[i * e_size + f] = inv_buf[f];
        }
        // 拷贝下三角部分
        for (idx_t j = 0; j < rnz; j++) {
            U.col_idx[U.row_ptr[i] + j] = A.col_idx[A.row_ptr[i] + lnz + 1 + j];
        }
        for (idx_t j = 0; j < rnz * e_size; j++) {
            U.vals[U.row_ptr[i] * e_size + j] = A.vals[(A.row_ptr[i] + lnz + 1) * e_size + j];
        }
    }
}


template class PointGS<int, double, double, double, 1>;
template class PointGS<int, double, double, double, 2>;
template class PointGS<int, double, double, double, 3>;
template class PointGS<int, double, double, double, 4>;

template class PointGS<int, float , double, double, 1>;
template class PointGS<int, float , double, double, 2>;
template class PointGS<int, float , double, double, 3>;
template class PointGS<int, float , double, double, 4>;

template class PointGS<int, float , double, float , 1>;
template class PointGS<int, float , double, float , 2>;
template class PointGS<int, float , double, float , 3>;
template class PointGS<int, float , double, float , 4>;

template class PointGS<int, float , float , float , 1>;
template class PointGS<int, float , float , float , 2>;
template class PointGS<int, float , float , float , 3>;
template class PointGS<int, float , float , float , 4>;
#ifdef USE_FP16
template class PointGS<int, __fp16, double, float , 1>;
template class PointGS<int, __fp16, double, float , 2>;
template class PointGS<int, __fp16, double, float , 3>;
template class PointGS<int, __fp16, double, float , 4>;

template class PointGS<int, __fp16, float , float , 1>;
template class PointGS<int, __fp16, float , float , 2>;
template class PointGS<int, __fp16, float , float , 3>;
template class PointGS<int, __fp16, float , float , 4>;
#endif